import modern_robotics as mr
import numpy as np

from .robot import Robot
from ..utils import TransformationUtils


class Frankie(Robot):
    def __init__(self):
        super().__init__()
        self._dof = 9
        self._q = np.zeros(self._dof)
        self._w_array = np.zeros((3, self._dof))
        self._v_array = np.zeros((3, self._dof))
        self._q_array = np.zeros((3, self._dof))

        self._w_array[:, 0] = [0.0, 0.0, 1.0]
        self._w_array[:, 1] = [0.0, 0.0, 0.0]
        self._w_array[:, 2] = [0.0, 0.0, 1.0]
        self._w_array[:, 3] = [0.0, 1.0, 0.0]
        self._w_array[:, 4] = [0.0, 0.0, 1.0]
        self._w_array[:, 5] = [0.0, -1.0, 0.0]
        self._w_array[:, 6] = [0.0, 0.0, 1.0]
        self._w_array[:, 7] = [0.0, -1.0, 0.0]
        self._w_array[:, 8] = [0.0, 0.0, -1.0]

        bx = 0.15
        bz = 0.38
        l1 = 0.333
        l2 = 0.316
        l3 = 0.384
        l4 = 0.107
        d1 = 0.0825
        d2 = 0.088

        self._q_array[:, 0] = [0.0, 0.0, 0.0]
        self._q_array[:, 1] = [0.0, 0.0, 0.0]
        self._q_array[:, 2] = [bx, 0.0, bz + l1]
        self._q_array[:, 3] = [bx, 0.0, bz + l1]
        self._q_array[:, 4] = [bx, 0.0, bz + l1 + l2]
        self._q_array[:, 5] = [bx + d1, 0.0, bz + l1 + l2]
        self._q_array[:, 6] = [bx, 0.0, bz + l1 + l2 + l3]
        self._q_array[:, 7] = [bx, 0.0, bz + l1 + l2 + l3]
        self._q_array[:, 8] = [bx + d2, 0.0, bz + l1 + l2 + l3]

        for i in range(self._dof):
            self._v_array[:, i] = np.cross(self._q_array[:, i], self._w_array[:, i])
        self._v_array[:, 1] = [1.0, 0.0, 0.0]

        self._Tbe = TransformationUtils.Txyz(bx + d2, 0.0, bz + l1 + l2 + l3 - l4) @ TransformationUtils.RPY(np.pi, 0.0,
                                                                                                             np.pi / 4)
        self._Tgb = np.eye(4)
        self._Tet = TransformationUtils.Txyz(0.0, 0.0, 0.1034)
        self._M = self._Tgb @ self._Tbe @ self._Tet
        self._S_array = mr.Adjoint(self._Tgb) @ np.vstack((self._w_array, self._v_array))
        self._B_array = mr.Adjoint(mr.TransInv(self._M)) @ self._S_array

        M0 = np.eye(4)
        M1 = np.eye(4)
        M2 = np.eye(4)
        M3 = TransformationUtils.Txyz(bx, 0.0, bz + l1)
        M4 = TransformationUtils.Txyz(bx, 0.0, bz + l1) @ TransformationUtils.Rx(-np.pi / 2)
        M5 = TransformationUtils.Txyz(bx, 0.0, bz + l1 + l2)
        M6 = TransformationUtils.Txyz(bx + d1, 0.0, bz + l1 + l2) @ TransformationUtils.Rx(np.pi / 2)
        M7 = TransformationUtils.Txyz(bx, 0.0, bz + l1 + l2 + l3)
        M8 = TransformationUtils.Txyz(bx, 0.0, bz + l1 + l2 + l3) @ TransformationUtils.Rx(np.pi / 2)
        M9 = TransformationUtils.Txyz(bx + d2, 0.0, bz + l1 + l2 + l3) @ TransformationUtils.Rx(np.pi)

        M10 = TransformationUtils.Txyz(bx + d2, 0.0, bz + l1 + l2 + l3 - l4) @ TransformationUtils.Rx(np.pi)
        M11 = TransformationUtils.Txyz(bx + d2, 0.0, bz + l1 + l2 + l3 - l4) @ TransformationUtils.RPY(np.pi, 0.0,
                                                                                                       np.pi / 4)
        M12 = TransformationUtils.Txyz(bx + d2, 0.0, bz + l1 + l2 + l3 - l4 - 0.1034) @ TransformationUtils.RPY(np.pi,
                                                                                                                0.0,
                                                                                                                np.pi / 4)

        self._Ms = [M0, M1, M2, M3, M4, M5, M6, M7, M8, M9, M10, M11, M12]

        self._q_lim = np.array([
            [-999999.0, -999999.0, -2.8973, -1.7628, -2.8973, -3.0718, -2.8973, -0.0175, -2.8973],
            [999999.0, 999999.0, 2.8973, 1.7628, 2.8973, -0.0698, 2.8973, 3.7525, 2.8973]
        ])
        self._qd_lim = np.array(
            [4.0, 4.0, 2.1750, 2.1750, 2.1750, 2.1750, 2.6100, 2.6100, 2.6100, 3.0, 3.0]
        )
